# 实现带时间的ns方程
# He XiaoMing 7.5.2 Example 2
from mpi4py import MPI
import numpy as np
from scipy.sparse import linalg

import ufl
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, la, log, mesh
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    form,
    functionspace,
)
from petsc4py import PETSc
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
from ufl import div, dx, grad, inner, TestFunctions, TrialFunctions, dot, nabla_grad, SpatialCoordinate
from dolfinx.nls.petsc import NewtonSolver
import GetConveregenceRate as gc



def true_velocity(x, t):
    return [u_x(x, t), u_y(x, t)]


def true_pressure(x, t):
    return pressure(x, t)


def u_x(x, t):
    """
    x方向速度的真实解
    """
    return (x[0]**2 * x[1]**2 + np.exp(-x[1])) * np.cos(2 * np.pi * t)


def u_y(x, t):
    """
    y方向速度的真实解
    """
    return (-2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * t)


def u_init(x):
    """
    返回u: IC
    """
    t1 = (x[0]**2 * x[1]**2 + np.exp(-x[1]))
    t2 = -2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])
    return [t1, t2]


def p_init_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1])


def p_end_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * T)


def pressure(x, t):
    """
    压强的真实解
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)


# ufl格式
def f_x(x, t):
    """
    f的x分量
    """
    term1 = -2 * ufl.pi * (x[0]**2 * x[1]**2 +
                           ufl.exp(-x[1])) * ufl.sin(2 * ufl.pi * t)
    # 非线性项
    nonlinear_term = 2 * x[0] * x[1]**2 * (x[0]**2 * x[1]**2 + ufl.exp(-x[1])) * ufl.cos(2 * ufl.pi * t)**2 + (
        2*x[0]**2 * x[1] - ufl.exp(-x[1])) * (-2/3*x[0]*x[1]**3 - ufl.pi*ufl.sin(ufl.pi*x[0])+2) * ufl.cos(2 * ufl.pi * t)**2
    term2 = (
        -2 * nu * x[0]**2
        - 2 * nu * x[1]**2
        - nu * ufl.exp(-x[1])
        + ufl.pi**2 * ufl.cos(ufl.pi * x[0]) * ufl.cos(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    return term1 + term2 + nonlinear_term


# He给出的数值算例不正确 第二项少了cos(2*pi*t)
def f_y(x, t):
    """
    f的y分量
    """
    term1 = -2 * ufl.pi * (-2/3 * x[0] * x[1]**3 + 2 - ufl.pi *
                           ufl.sin(ufl.pi * x[0])) * ufl.sin(2 * ufl.pi * t)
    term2 = (
        4 * nu * x[0] * x[1]
        - nu * ufl.pi**3 * ufl.sin(ufl.pi * x[0])
        + 2 * ufl.pi * (2 - ufl.pi * ufl.sin(ufl.pi *
                        x[0])) * ufl.sin(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    #  非线性项
    nonlinear_term = -2*x[0]*x[1]**2 * (-2/3*x[0]*x[1]**3 - ufl.pi*ufl.sin(ufl.pi*x[0])+2) * ufl.cos(2 * ufl.pi * t)**2 + \
        (-2/3*x[1]**3-ufl.pi**2*ufl.cos(ufl.pi*x[0])) * (x[0] **
                                                         2 * x[1]**2 + ufl.exp(-x[1])) * ufl.cos(2 * ufl.pi * t)**2
    return term1 + term2 + nonlinear_term


def f(x, t):
    """
    返回f(list)
    """
    return [f_x(x, t), f_y(x, t)]


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def get_norm(t: float, save_=True):
    """
    传入时间t
    """
    # 计算误差
    u_true = Function(Q)
    p_true = Function(S)

    u_true.interpolate(lambda x: np.array([u_x(x, t), u_y(x, t)]))

    def p_current(x):  # 只能这样设置时间 比较丑陋
        return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)
    p_true.interpolate(p_current)

    u_error = Function(Q)
    u_error.x.array[:] = u_h.x.array - u_true.x.array
    p_error = Function(S)
    p_error.x.array[:] = p_h.x.array - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((grad(u_error)**2 + u_error**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((grad(p_error)**2 + p_error**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)
    data_ = [
        n_list[idx],
        u_h1_E, p_h1_E,
        u_l2_E, p_l2_E,
        u_linf_E, p_linf_E,
    ]
    if MPI.COMM_WORLD.rank == 0:
        if save_:  # 是否写入为excel
            gc.save_data(FILENAME, data_)
        print(
            f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.")
        print(
            f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.")
        print(
            f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.")


type_ = "newton"  # Available: "stokes" "picard" "newton"
assert type_ in ['stokes', 'picard',
                 'newton'], "以下为有效迭代类型: ['stokes', 'picard', 'newton']"
FILENAME = f"ns-{type_}.xlsx"
gc.create_file(FILENAME)
# Params
n_list = [
    [8, 2],
    [16, 4],
    [32, 8],
    [64, 16],
    [128, 32],
]
rtol = 1e-8
max_iter = 100
idx = int(input("请输入网格划分的索引: "))
nu = 1
t = 0  # 初始时刻
T = 1  # 终止时刻
time_steps = n_list[idx][0] ** 2  # 时间步 取得细致一点 否则误差无法减小...
dt = T / time_steps

msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(
        [1, 0])], n=n_list[idx], cell_type=CellType.triangle
)
x = SpatialCoordinate(msh)  # 定义空间坐标
print(f"Mesh = {n_list[idx]}")

P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
mel = mixed_element([P2, P1])  # 定义混合元空间 分别对应速度 压力
ME = functionspace(msh, mel)  # 实例化混合元空间
sol = Function(ME)  # 定义解函数
sol.x.array[:] = 0.0  # 随机初始化
Q, _ = ME.sub(0).collapse()  # 速度空间
S, _ = ME.sub(1).collapse()  # 压力空间
u_bc = Function(Q)  # 真实速度函数空间
p_bc = Function(S)  # 真实压力函数空间

# 标记边界
facets = locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(
        x[0], 1.0) | np.isclose(x[1], -0.25) | np.isclose(x[1], 0.0),
)
v_dofs = fem.locate_dofs_topological(
    (ME.sub(0), Q), entity_dim=1, entities=facets)

# 固定压力


def marker(x): return np.isclose(x[0], 0.0) & np.isclose(x[1], -0.25)


p_dofs = fem.locate_dofs_geometrical((ME.sub(1), S), marker)

u_n = Function(Q)  # 上一时间步长的计算结果
p_n = Function(S)
u_n.interpolate(u_init)
p_n.interpolate(p_init_expr)
u_old = Function(Q)  # 上一步的迭代结果
u_old.x.array[:] = 0.0


for time_step in range(time_steps+1):
    # Define trial and test functions
    u, p = TrialFunctions(ME)  # trial: u, p; test: v, q
    v, q = TestFunctions(ME)  # Test Functions
    f_expr = ufl.as_vector([f_x(x, t), f_y(x, t)])  # 调整f为ufl格式

    # BCs
    u_bc.interpolate(lambda x: true_velocity(x, t))
    p_bc.interpolate(lambda x: pressure(x, t))
    bc_vel = dirichletbc(u_bc, v_dofs, ME.sub(0))  # 速度边界条件
    bc_pressure = dirichletbc(p_bc, p_dofs, ME.sub(1))
    bcs = [bc_vel, bc_pressure]

    # Weak Formulation
    a = 2 * nu * inner(D(u), D(v)) * dx
    a += inner(div(u), q) * dx
    a -= inner(div(v), p) * dx
    a += inner(u, v) / dt * dx
    f_expr = ufl.as_vector([f_x(x, t), f_y(x, t)])  # 调整f为ufl格式
    L = inner(f_expr, v) * dx  # 右端项
    L += inner(u_n, v) / dt * dx
    # 迭代
    for i in range(max_iter):
        # 共用
        if type_ == "stokes":
            L = inner(f_expr, v) * dx + inner(u_n, v) / dt * dx  # 右端项
            L -= inner(dot(u_old, nabla_grad(u_old)), v) * dx
        elif type_ == "picard":
            a = 2 * nu * inner(D(u), D(v)) * dx
            a += inner(div(u), q) * dx
            a -= inner(div(v), p) * dx
            a += inner(dot(u_old, nabla_grad(u)), v) * dx
            a += inner(u, v) / dt * dx
        elif type_ == "newton":
            a = 2 * nu * inner(D(u), D(v)) * dx
            a += inner(div(u), q) * dx
            a -= inner(div(v), p) * dx
            a += inner(dot(u_old, nabla_grad(u)), v) * dx
            a += inner(dot(u, nabla_grad(u_old)), v) * dx
            a += inner(u, v) / dt * dx
            L = inner(f_expr, v) * dx  # 右端项
            L += inner(dot(u_old, nabla_grad(u_old)), v) * dx
            L += inner(u_n, v) / dt * dx

        a_compiled = fem.form(a)
        L_compiled = fem.form(L)
        A = fem.create_matrix(a_compiled)
        b = fem.create_vector(L_compiled)
        A_scipy = A.to_scipy()
        fem.assemble_matrix(A, a_compiled, bcs=bcs)
        fem.assemble_vector(b.array, L_compiled)
        fem.apply_lifting(b.array, [a_compiled], [bcs])
        b.scatter_reverse(la.InsertMode.add)
        [bc.set(b.array) for bc in bcs]
        A_inv = linalg.splu(A_scipy)
        sol.x.array[:] = A_inv.solve(b.array)  # 怎么设置循环？
        u_old_array = u_old.x.array.copy()
        u_old.x.array[:] = sol.sub(0).collapse().x.array
        realative_err = np.linalg.norm(
            u_old.x.array - u_old_array) / (np.linalg.norm(u_old_array) + 1e-10)
        if realative_err < rtol:
            break

    u_h = sol.sub(0).collapse()
    p_h = sol.sub(1).collapse()
    u_n.x.array[:] = u_h.x.array  # 从当前解中提取速度
    p_n.x.array[:] = p_h.x.array  # 从当前解中提取压力
    u_n.x.scatter_forward()
    p_n.x.scatter_forward()
    print(f"t = {t}, Iter = {i} finish.")
    t += dt
get_norm(t=T)  # 计算误差

# preview
# Mesh = [8, 2]
# H1 norm for velocity: 0.003638324959338088, pressure: 0.3759448936133618.
# L2 norm for velocity: 9.027566986538663e-05, pressure: 0.05907687028871413.
# Linf norm for velocity: 0.0006696483617828486, pressure: 0.4271632240023071.
# Mesh = [16, 4]
# H1 norm for velocity: 0.0007026378250481561, pressure: 0.09497213265296307.
# L2 norm for velocity: 1.9123126569176187e-05, pressure: 0.018809299848414496.
# Linf norm for velocity: 9.107936769092717e-05, pressure: 0.11644229832188069.
